#include "txr_algo_dlm_rec_demo.h"
#include <QFileDialog>
#include <QTextStream>
#include <QDir>
#include <QFile>
#include <sstream>
#include <fstream>
#include <iomanip> 
#include "../txr_algo_dlm_rec/txr_algo_dlm_rec.h"
#include "../tk_utility/tk_timer.hpp"
#include <QThread>
using namespace txr_algo_dlm_rec;

#include <opencv2\highgui.hpp>
#include <opencv2\imgproc.hpp>

// Global map to store label mappings
std::map<int, std::string> g_label_map;

/**
 * Read COCO format label file and create label mapping
 * @param fileName Path to the label file
 * @return Map of label index to label string
 */
std::map<int, std::string> readCOCOLabel(const std::string& fileName)
{
	std::map<int, std::string> coco_label;
	std::ifstream file(fileName);
	if (!file.is_open())
	{
		std::cout << "read file error: " << fileName << std::endl;
	}
	std::string strLine;
	int index = 0;
	while (getline(file, strLine))
	{
		coco_label.insert({ index, strLine });
		index++;
	}
	file.close();
	return coco_label;
}

/**
 * Read and preprocess image file into float array
 * @param file_path Path to image file
 * @param v_data Output vector to store preprocessed image data
 * @param row Output parameter for image height
 * @param col Output parameter for image width
 * @return True if successful, false otherwise
 */
bool ReadImage(QString file_path, QVector<float> &v_data, int & row, int &col)
{
	using namespace cv;
	Mat src = imread(file_path.toLocal8Bit().data(), IMREAD_ANYCOLOR | IMREAD_ANYDEPTH);
	if (src.empty())
	{
		return false;
	}
	if (src.channels() != 3 && src.channels() != 1)
	{
		return false;
	}
	row = src.rows;
	col = src.cols;

	v_data.resize(row * col * 3);

	int img_size = row*col;
	if (src.channels() == 3)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3w color_value = src.ptr<Vec3w>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3b color_value = src.ptr<Vec3b>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}
	if (src.channels() == 1)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<ushort, 1> color_value = src.ptr<Vec<ushort, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<uchar, 1> color_value = src.ptr<Vec<uchar, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}

	return true;
}

/**
 * Draw detection results on image and save
 * @param unit Detection unit containing image data
 * @param dest_path Output image path
 * @param result Detection results to visualize
 */
void MarkImage(st_detect_unit unit, QString dest_path, st_result_info & result)
{
	st_img_rgb img = unit.img;
	using namespace cv;
	Mat mat;
	mat.create(img.h, img.w, CV_8UC3);
	int img_size = img.h*img.w;
	for (int i = 0; i < img.h; i++)
	{
		for (int j = 0; j < img.w; j++)
		{
			mat.at<Vec3b>(i, j)[2] = img.pp_rgbf[0][i*img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[1] = img.pp_rgbf[1][i*img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[0] = img.pp_rgbf[2][i*img.w + j] * 255.f;
		}
	}
	
	std::stringstream stream;
	int code_num = result.vaild_num;
	for (int n = 0; n < code_num; ++n)
	{
		int label = result.code[n].code_label;
		float prob = result.code[n].prob;
		std::string code = g_label_map.find(label) != g_label_map.end() ? g_label_map[label] : " ";

		stream << "[" << code << "]:" << std::fixed << std::setprecision(2) << prob << " ";
	}
	cv::Point pt(0, 20);
	putText(mat, stream.str(), pt, 0, 0.3, cv::Scalar(0, 0, 255), 0.3);
	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}
	imwrite(dest_path.toLocal8Bit().data(), mat);
}

/**
 * Constructor for main window
 * @param parent Parent widget
 */
txr_algo_dlm_rec_demo::txr_algo_dlm_rec_demo(QWidget *parent)
    : QMainWindow(parent)
{
    ui.setupUi(this);
	connect(ui.pushButtonOpenTestDir, SIGNAL(clicked()), this, SLOT(OnOpenTest()));
	connect(ui.pushButtonOpenModelDir, SIGNAL(clicked()), this, SLOT(OnOpenModel()));
	connect(ui.pushButtonTest, SIGNAL(clicked()), this, SLOT(OnTest()));
	ui.progressBar->setRange(0, 100);
	ui.progressBar->setValue(0);

	g_label_map = readCOCOLabel("rec.names");

}

/**
 * Open test directory dialog handler
 */
void txr_algo_dlm_rec_demo::OnOpenTest()
{
	QString dir = QFileDialog::getExistingDirectory(this);
	if (dir.isEmpty())
	{
		return;
	}
	ui.labelDirTest->setText(dir);
}

/**
 * Open model directory dialog handler
 */
void txr_algo_dlm_rec_demo::OnOpenModel()
{
	QString dir_path = QFileDialog::getExistingDirectory(this);
	if (dir_path.isEmpty())
	{
		return;
	}
	ui.labelDirModel->setText(dir_path);

	QDir dir(dir_path);
	QStringList filters;
	filters << "*.rec1" << "*.orec1";

	QStringList files = dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	ui.comboBoxModel->clear();
	ui.comboBoxModel->addItems(files);
}

/**
 * Test button click handler
 */
void txr_algo_dlm_rec_demo::OnTest()
{
	st_test_info info;

	info.test_dir = ui.labelDirTest->text();

	info.model_path = ui.labelDirModel->text() + "\\" + ui.comboBoxModel->currentText();

	info.save_img = ui.checkBoxSaveImg->isChecked();

	DlmTest(info);
}

/**
 * Run DLM test on images
 * @param info Test configuration info
 */
void txr_algo_dlm_rec_demo::DlmTest(st_test_info &info)
{
	TxrAlgoDlmRec dlm;
	{
		st_dlm_cfg cfg;

		QFileInfo model_file = QFileInfo(info.model_path);
		QString model_name = model_file.absoluteDir().path() + "\\" + model_file.baseName();
		QString model_suffix = model_file.suffix();

		sprintf_s(cfg.file_path_name, MAX_PATH_BYTES, "%s", model_name.toLocal8Bit().data());
		if (model_suffix == "rec1")
		{
			cfg.type = DLM_REC_TYPE_V1;
		}
		else if (model_suffix == "orec1")
		{
			cfg.type = DLM_REC_TYPE_V1_O;
		}
		if (!dlm.Load(cfg))
		{
			ui.plainTextEditInfo->appendPlainText("模型加载失败!");
			return;
		}
	}
	int batch_size = dlm.BatchSize();

	

	//============================================================
	{
		QDir dir(info.test_dir + "\\result");
		dir.removeRecursively();
	}
	//============================================================
	ui.plainTextEditInfo->clear();

	QStringList test_files;
	{
		QDir dir(info.test_dir);
		QStringList filters;
		filters << "*.png" << "*.bmp"<<"*.jpg";
		test_files =
			dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	}

	int file_count = 0;
	ui.progressBar->setValue(0);

	QVector<QString> v_img_path;
	QVector<QVector<float>> vv_data;
	QVector<txr_algo_dlm_rec::st_detect_unit> v_unit;
	v_img_path.resize(batch_size);
	vv_data.resize(batch_size);
	v_unit.resize(batch_size);

	// 用于记录当前批次中已处理的图像数量
	int batch_count = 0;
	for each (QString file in test_files)
	{
		// 打印批处理大小
		printf("batch size: %d\n", batch_size);
		// 如果当前批次未满,继续添加图像
		if (batch_count<batch_size)
		{
			// 获取当前批次的图像路径引用
			QString & img_path = v_img_path[batch_count];
			img_path = info.test_dir + "\\" + file;
			
			// 获取当前批次的图像数据和检测单元引用
			QVector<float> & v_data = vv_data[batch_count];
			txr_algo_dlm_rec::st_detect_unit & unit = v_unit[batch_count];
			
			// 读取图像,获取行列数
			int row, col;
			if (!ReadImage(img_path, v_data, row, col))
			{
				// 读取失败则跳过当前图像
				continue;
			}

			// 设置检测单元的图像信息
			unit.img.w = col;
			unit.img.h = row;
			unit.img.c = 3; // RGB三通道
			
			// 设置RGB三个通道的数据指针
			// 注意:这里假设图像数据是按照RRRGGGBBB的方式存储的
			unit.img.pp_rgbf[0] = v_data.data() + col*row * 0; // R通道起始位置
			unit.img.pp_rgbf[1] = v_data.data() + col*row * 1; // G通道起始位置  
			unit.img.pp_rgbf[2] = v_data.data() + col*row * 2; // B通道起始位置
			batch_count++;
		}

		// 如果批次未满且不是最后一个文件,继续下一个循环
		if (batch_count != batch_size && file != test_files.last())
		{
			continue;
		}

		// 执行批量检测
		{
			printf("batch_count: %d\n", batch_count); // 打印当前批次大小用于调试
			tk::timer::universal_short_timer timer;
			// 注意:这里传入batch_count而不是batch_size,因为最后一批可能不满
			dlm.Detect(v_unit.data(), batch_count);
			printf("dlm detect openvino use time: %f ms \n", timer.elapsed_ms());
		}
		for (int i = 0; i < batch_count; ++i)
		{
			st_result_info result = v_unit[i].result;

			int code_num = result.vaild_num;
			std::string str_info;
			for (int n = 0; n < code_num; ++n)
			{
				int label = result.code[n].code_label;
				float prob = result.code[n].prob;
				std::string code = g_label_map.find(label) != g_label_map.end() ? g_label_map[label] : " ";

				str_info += "[" + code + "]:" + std::to_string(prob) + "\n";
			}

			QFileInfo file_info = QFileInfo(v_img_path[i]);
			QString mark_file = info.test_dir + QString("\\result\\") + file_info.baseName() + ".png";
			MarkImage(v_unit[i], mark_file, result);

			QString text = QString("no:%1 file:%2\n%3")
				.arg(file_count)
				.arg(file_info.baseName())
				.arg(str_info.c_str());
			ui.plainTextEditInfo->insertPlainText(text);
			file_count++;
		}
		batch_count = 0;

		ui.progressBar->setValue(file_count * 100 / test_files.size());
		QCoreApplication::processEvents();
	}
	ui.progressBar->setValue(100);
	printf("done!==========================================================\n");
}